| from ultralytics import YOLO |
| import numpy as np |
| from PIL import Image |
| import io |
| import base64 |
| import torch |
|
|
| def hf_predict(image_bytes): |
| """ |
| Inference function for Hugging Face Inference API |
| |
| Args: |
| image_bytes: Image file bytes from a POST request |
| |
| Returns: |
| dict: Detection results in a format compatible with the Inference API |
| """ |
| |
| image = Image.open(io.BytesIO(image_bytes)).convert('RGB') |
| |
| |
| model = YOLO('best.pt') |
| |
| |
| results = model(image) |
| |
| |
| result = results[0] |
| boxes = result.boxes.xyxy.cpu().numpy() |
| scores = result.boxes.conf.cpu().numpy() |
| class_ids = result.boxes.cls.cpu().numpy() |
| |
| |
| class_names = ["card"] |
| labels = [class_names[int(i)] for i in class_ids] |
| |
| |
| detections = [] |
| for box, score, label in zip(boxes, scores, labels): |
| x1, y1, x2, y2 = box |
| detections.append({ |
| "box": { |
| "xmin": int(x1), |
| "ymin": int(y1), |
| "xmax": int(x2), |
| "ymax": int(y2) |
| }, |
| "score": float(score), |
| "label": label |
| }) |
| |
| return {"detections": detections} |
|
|
| def visualize(image_bytes, detections): |
| """ |
| Visualize the detections on the image |
| |
| Args: |
| image_bytes: Image file bytes |
| detections: Detection results |
| |
| Returns: |
| str: Base64 encoded image with visualizations |
| """ |
| try: |
| import cv2 |
| |
| |
| nparr = np.frombuffer(image_bytes, np.uint8) |
| image = cv2.imdecode(nparr, cv2.IMREAD_COLOR) |
| |
| |
| for det in detections["detections"]: |
| box = det["box"] |
| score = det["score"] |
| label = det["label"] |
| |
| x1, y1 = box["xmin"], box["ymin"] |
| x2, y2 = box["xmax"], box["ymax"] |
| |
| cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) |
| cv2.putText(image, f"{label}: {score:.2f}", (x1, y1 - 10), |
| cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) |
| |
| |
| _, buffer = cv2.imencode('.jpg', image) |
| return base64.b64encode(buffer).decode('utf-8') |
| except: |
| return None |
|
|
