add demo

.gitignore

@@ -49,4 +49,7 @@ Thumbs.db
 
 # --- Log Files ---
 # Ignore log files, which can become large and are specific to a run.
-*.log
+*.log
+*.jpg
+*.jpeg
+*.png

app.py

@@ -1,57 +1,184 @@
+# app.py
+import io
+import uvicorn
 import numpy as np
+from PIL import Image, ImageDraw, ImageFont
+from typing import List
+from fastapi import FastAPI, UploadFile, File, Request, Form
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+from fastapi.templating import Jinja2Templates
+import onnxruntime as ort
 import cv2
-import gradio as gr
-from ultralytics import YOLO
-import os
 from huggingface_hub import hf_hub_download
-from dotenv import load_dotenv
+import os
+import uuid
+
+# --- FastAPI and Template Setup ---
+app = FastAPI(title="YOLOv8 ONNX Object Detection Demo")
+
+# Mount a static directory to serve saved images
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+templates = Jinja2Templates(directory="templates")
+
+# --- Model Loading and Configuration ---
+# Download the ONNX model file and get its path
+try:
+    onnx_model_path = hf_hub_download(repo_id="tententgc/Iskyn", filename="best.onnx")
+    session = ort.InferenceSession(onnx_model_path)
+    print("ONNX model loaded successfully.")
+except Exception as e:
+    print(f"Failed to load ONNX model: {e}")
+    session = None
+
+if session:
+    input_name = session.get_inputs()[0].name
+    output_names = [output.name for output in session.get_outputs()]
+    input_shape = session.get_inputs()[0].shape[2:]  # Get the expected image size
+else:
+    input_name = None
+    output_names = []
+    input_shape = (640, 640)  # Default size if model fails to load
+
+# Define the class names for your model
+# IMPORTANT: Update this with the actual class names your model was trained on
+CLASSES = [
+    "melasma", "acne", "wrinkle" 
+]
+
+# A dictionary to map class names to colors for plotting
+COLORS = {
+    "melasma": "red",
+    "acne": "green",
+    "wrinkle": "blue",
+    # Add more classes and colors as needed
+}
+
+# --- Helper Functions ---
+def preprocess_image(image: Image.Image, size: tuple) -> np.ndarray:
+    """Preprocesses an image for model inference."""
+    image = image.resize(size)
+    image = np.array(image)
+    image = image.transpose(2, 0, 1)  # HWC to CHW
+    image = np.expand_dims(image, axis=0) # Add batch dimension
+    image = image.astype(np.float32) / 255.0  # Normalize
+    return image
+
+def postprocess_output(output, original_size, input_shape, conf_threshold=0.25, iou_threshold=0.45):
+    """Post-processes the model output to get bounding boxes, scores, and class IDs."""
+    output = np.squeeze(output).T
+    scores = np.max(output[:, 4:], axis=1)
+    filtered_indices = scores > conf_threshold
+    output = output[filtered_indices]
+    scores = scores[filtered_indices]
+
+    if not len(output):
+        return []
 
+    boxes = output[:, :4]
+    boxes[:, 0] -= boxes[:, 2] / 2
+    boxes[:, 1] -= boxes[:, 3] / 2
+    boxes[:, 2] += boxes[:, 0]
+    boxes[:, 3] += boxes[:, 1]
 
-print("Loading environment variables...")
-load_dotenv()
+    class_ids = np.argmax(output[:, 4:], axis=1)
+    indices = cv2.dnn.NMSBoxes(boxes.astype(np.int32), scores.astype(np.float32), conf_threshold, iou_threshold)
 
-HF_REPO_ID = "tententgc/Iskyn"
-MODEL_FILENAME = "best.onnx"
+    detections = []
+    if len(indices) > 0:
+        for i in indices.flatten():
+            box = boxes[i]
+            x1, y1, x2, y2 = box.astype(int)
+            class_id = class_ids[i]
+            score = scores[i]
 
-print(f"Downloading '{MODEL_FILENAME}' from '{HF_REPO_ID}'...")
+            original_width, original_height = original_size
+            resized_width, resized_height = input_shape
+            x1 = int(x1 * original_width / resized_width)
+            y1 = int(y1 * original_height / resized_height)
+            x2 = int(x2 * original_width / resized_width)
+            y2 = int(y2 * original_height / resized_height)
 
-model_path = hf_hub_download(
-    repo_id=HF_REPO_ID,
-    filename=MODEL_FILENAME,
-    token=os.getenv("HF_TOKEN") 
-)
+            detections.append({
+                "class_name": CLASSES[class_id],
+                "confidence": float(score),
+                "box": [x1, y1, x2, y2]
+            })
+    return detections
 
-print(f"Model downloaded to: {model_path}")
+def draw_boxes_on_image(image, detections):
+    """Draws bounding boxes, class names, and confidence scores on an image."""
+    draw = ImageDraw.Draw(image)
+    try:
+        font = ImageFont.truetype("arial.ttf", 30)
+    except IOError:
+        font = ImageFont.load_default()
+        print("Arial font not found, using default font.")
 
+    for detection in detections:
+        box = detection['box']
+        class_name = detection['class_name']
+        confidence = detection['confidence']
+        
+        color = COLORS.get(class_name, "white")
+        draw.rectangle(box, outline=color, width=3)
+        
+        label = f"{class_name}: {confidence:.2f}"
+        
+        # Use textbbox() to get text dimensions
+        text_x, text_y, text_width, text_height = draw.textbbox((0, 0), label, font=font)
+        
+        # Position text slightly above the top-left corner
+        text_position_y = box[1] - text_height - 5
+        if text_position_y < 0:
+            text_position_y = box[1] + 5 # Draw below if not enough space above
+            
+        draw.rectangle([box[0], text_position_y, box[0] + text_width, text_position_y + text_height], fill=color)
+        draw.text((box[0], text_position_y), label, fill="black", font=font)
+    return image
 
+# --- FastAPI Endpoints ---
+@app.get("/", response_class=HTMLResponse)
+async def read_root(request: Request):
+    """Serve the HTML interface."""
+    return templates.TemplateResponse("index.html", {"request": request, "image_url": None, "error_message": None})
 
-print("Loading YOLO model...")
-onnx_model = YOLO(model_path) 
-print("Model loaded successfully.")
+@app.post("/predict_web", response_class=HTMLResponse)
+async def predict_web(request: Request, file: UploadFile = File(...)):
+    """Handle image upload, run detection, and return plotted image."""
+    if not session:
+        return templates.TemplateResponse("index.html", {"request": request, "error_message": "ONNX model not loaded."})
+    
+    if not file.content_type.startswith("image/"):
+        return templates.TemplateResponse("index.html", {"request": request, "error_message": "Invalid file type. Please upload an image."})
 
+    try:
+        image_data = await file.read()
+        image = Image.open(io.BytesIO(image_data)).convert("RGB")
+        original_size = image.size
 
+        # Preprocess, run inference, and post-process
+        preprocessed_image = preprocess_image(image, size=input_shape)
+        outputs = session.run(output_names, {input_name: preprocessed_image})
+        detections = postprocess_output(outputs, original_size, input_shape)
 
-def predict_image(image_filepath, conf_threshold, iou_threshold):
-    results = onnx_model.predict(
-        image_filepath,
-        conf=conf_threshold,
-        iou=iou_threshold
-    )
-    result = results[0]
-    im_array = result.plot()
-    im_rgb = cv2.cvtColor(im_array, cv2.COLOR_BGR2RGB)
-    return im_rgb
+        # Draw boxes on the original image
+        plotted_image = draw_boxes_on_image(image.copy(), detections)
 
-iface = gr.Interface(
-    fn=predict_image,
-    inputs=[
-        gr.Image(type="filepath", label="Upload Image"),
-        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence Threshold"),
-        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU Threshold")
-    ],
-    outputs=gr.Image(type="numpy", label="Result"),
-    title="Detection Face Skin",
-    description="Upload an image and adjust the thresholds to fine-tune detection." 
-)
+        # Create a unique filename and save the plotted image
+        unique_filename = f"{uuid.uuid4()}.jpg"
+        output_image_path = os.path.join("static", "output", unique_filename)
+        plotted_image.save(output_image_path)
+        
+        image_url = f"/static/output/{unique_filename}"
+        
+        return templates.TemplateResponse("index.html", {"request": request, "image_url": image_url})
+    
+    except Exception as e:
+        return templates.TemplateResponse("index.html", {"request": request, "error_message": f"An error occurred: {e}"})
 
-iface.launch()
+if __name__ == "__main__":
+    # Create the static/output directory if it doesn't exist
+    os.makedirs(os.path.join("static", "output"), exist_ok=True)
+    uvicorn.run(app, host="127.0.0.1", port=8000)

requirements.txt

@@ -47,7 +47,7 @@ h11==0.16.0
 hf-xet==1.1.10
 httpcore==1.0.9
 httpx==0.28.1
-huggingface-hub==0.34.4
+huggingface-hub==0.35.0
 humanfriendly==10.0
 idna==3.10
 ifaddr==0.2.0

templates/index.html

@@ -0,0 +1,43 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>YOLO Object Detection Demo</title>
+    <style>
+        body { font-family: Arial, sans-serif; padding: 20px; text-align: center; }
+        h1 { color: #333; }
+        .container { max-width: 800px; margin: auto; }
+        form { margin-top: 20px; padding: 20px; border: 1px solid #ddd; border-radius: 8px; }
+        .image-display { margin-top: 20px; }
+        .image-display img { max-width: 100%; border: 1px solid #ccc; border-radius: 8px; }
+        .error-message { color: red; font-weight: bold; }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <h1>YOLO Object Detection Demo</h1>
+        <p>Upload an image to perform object detection.</p>
+
+        <form action="/predict_web" method="post" enctype="multipart/form-data">
+            <input type="file" name="file" accept="image/*" required>
+            <br><br>
+            <button type="submit">Upload and Detect</button>
+        </form>
+
+        {% if image_url %}
+        <div class="image-display">
+            <h2>Detection Result:</h2>
+            <img src="{{ image_url }}" alt="Detected Objects">
+        </div>
+        {% endif %}
+
+        {% if error_message %}
+        <div class="error-message">
+            <p>{{ error_message }}</p>
+        </div>
+        {% endif %}
+
+    </div>
+</body>
+</html>

test_api.py

@@ -0,0 +1,102 @@
+import requests
+from PIL import Image, ImageDraw, ImageFont
+import io
+import os
+
+# The URL of your FastAPI predict endpoint
+url = "http://127.0.0.1:8000/predict"
+
+
+image_path = "acne-face-2-18.jpg" 
+output_path = "result.jpg" 
+
+
+COLORS = {
+    "acne": "red",
+    "melasma": "green",
+    "wrinkle": "blue"
+}
+
+def draw_boxes_on_image(image, detections):
+    """Draws bounding boxes, class names, and confidence scores on an image."""
+    draw = ImageDraw.Draw(image)
+    try:
+        # Try to use a better-looking font if available
+        font = ImageFont.truetype("arial.ttf", 20)
+    except IOError:
+        font = ImageFont.load_default()
+        print("Arial font not found, using default font.")
+
+    for detection in detections:
+        box = detection['box']
+        class_name = detection['class_name']
+        confidence = detection['confidence']
+        
+        # Get color based on class name, defaulting to a solid color if not found
+        color = COLORS.get(class_name, "white")
+        
+        # Draw the rectangle
+        draw.rectangle(box, outline=color, width=3)
+        
+        # Create the label text with class name and confidence
+        label = f"{class_name}: {confidence:.2f}"
+        
+        # Use textbbox() to get text dimensions
+        # It returns a tuple: (left, top, right, bottom)
+        bbox = draw.textbbox((0, 0), label, font=font)
+        text_width = bbox[2] - bbox[0]
+        text_height = bbox[3] - bbox[1]
+        
+        # Define text position slightly above the top-left corner of the box
+        text_x = box[0]
+        text_y = box[1] - text_height - 5  # 5 pixels padding
+        
+        # Ensure text is not drawn off the top of the image
+        if text_y < 0:
+            text_y = box[1] + 5 # Draw below the box if no space above
+        
+        # Draw a filled background for the text for better visibility
+        draw.rectangle([text_x, text_y, text_x + text_width, text_y + text_height], fill=color)
+        
+        # Draw the label text
+        draw.text((text_x, text_y), label, fill="black", font=font)
+        
+    return image
+
+try:
+    # Check if the image file exists
+    if not os.path.exists(image_path):
+        raise FileNotFoundError(f"Error: The image file was not found at {image_path}")
+
+    # Open the image file in binary mode
+    with open(image_path, "rb") as f:
+        files = {"file": f}
+        
+        # Send the POST request to the FastAPI endpoint
+        response = requests.post(url, files=files)
+        
+    # Check for a successful response (status code 200)
+    if response.status_code == 200:
+        detections = response.json().get("detections", [])
+        
+        if detections:
+            print("Detections found:", detections)
+            # Load the original image again for plotting
+            original_image = Image.open(image_path).convert("RGB")
+            
+            # Draw the detections on the image
+            plotted_image = draw_boxes_on_image(original_image, detections)
+            
+            # Save the new image with the plots
+            plotted_image.save(output_path)
+            print(f"Success! Plotted image saved to: {output_path}")
+            
+        else:
+            print("No objects were detected.")
+            
+    else:
+        print(f"Error: API returned status code {response.status_code}")
+        print("Response:", response.text)
+
+except requests.exceptions.RequestException as e:
+    print(f"An error occurred while connecting to the API: {e}")