Update app.py

app.py

@@ -1,79 +1,61 @@
+import gradio as gr
 import os
 import subprocess
-from fastapi import FastAPI, UploadFile, File, HTTPException
-from fastapi.responses import FileResponse, PlainTextResponse
-import shutil
+import cv2
+import numpy as np
 
-app = FastAPI()
-
-# Define paths and parameters
-weights_path = 'yolo/yolov7-main/runs/train/best.pt'
-img_size = 640
-conf = 0.20
-source_folder = 'dataset/images/train/'
-output_folder = 'out/fixed_folder/'
-
-# Ensure folders exist
-os.makedirs(source_folder, exist_ok=True)
-os.makedirs(output_folder, exist_ok=True)
+# Define the detect function
+def detect_and_crop(input_image):
+    # Define paths and parameters
+    weights_path = 'yolo/yolov7-main/runs/train/best.pt'
+    img_size = 640
+    conf = 0.20
+    source = 'dataset/images/train/'  # Folder for input images
 
-# Root endpoint to prevent 404 on /
-@app.get("/")
-async def root():
-    return PlainTextResponse("Welcome to the YOLOv7 Object Detection API. Use the /detect endpoint to upload an image.")
+    # Ensure the input image folder exists
+    os.makedirs(source, exist_ok=True)
 
-# Optional: Add a route for favicon.ico to avoid 404
-@app.get("/favicon.ico", include_in_schema=False)
-async def favicon():
-    return PlainTextResponse("", status_code=204)
+    # Save the input image to the source directory
+    input_image.save(os.path.join(source, 'input_image.jpg'))
 
-# Define the detect function
-def detect_and_crop(image_path: str):
     # Run the detection command
     command = [
         'python', 'yolo/yolov7-main/detect.py',
         '--weights', weights_path,
         '--conf-thres', str(conf),
         '--img-size', str(img_size),
-        '--source', image_path,
+        '--source', os.path.join(source, 'input_image.jpg'),
         '--project', 'out/',  # Output directory
         '--name', 'fixed_folder',  # Folder name for results
         '--exist-ok'  # Don't increment folder name
     ]
 
-    # Execute the command and check for errors
-    result = subprocess.run(command, capture_output=True, text=True)
-    if result.returncode != 0:
-        raise HTTPException(status_code=500, detail=f"Detection failed: {result.stderr}")
+    # Execute the command
+    subprocess.run(command)
+
+    # Load the result image
+    output_image_path = 'out/fixed_folder/input_image_upscaled.jpg'
 
-    # Locate the output image in the expected output directory
-    output_files = os.listdir(output_folder)
-    output_image_path = None
-    for file_name in output_files:
-        if file_name.endswith(".jpg") or file_name.endswith(".jpeg") or file_name.endswith(".png"):
-            output_image_path = os.path.join(output_folder, file_name)
-            break
+    # Check if the image exists
+    if not os.path.exists(output_image_path):
+        return "No output image found."
 
-    if not output_image_path or not os.path.exists(output_image_path):
-        raise HTTPException(status_code=404, detail="Output image not found.")
+    # Read the output image
+    output_image = cv2.imread(output_image_path)
 
-    return output_image_path
+    # Convert BGR (OpenCV format) to RGB (Gradio format)
+    output_image = cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB)
 
-# FastAPI endpoint to accept an image, perform detection, and return the processed image
-@app.post("/detect")
-async def detect_endpoint(file: UploadFile = File(...)):
-    # Save the uploaded file to the source folder
-    input_image_path = os.path.join(source_folder, 'input_image.jpg')
-    with open(input_image_path, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
+    return output_image
 
-    # Perform detection and get the path to the output image
-    try:
-        output_image_path = detect_and_crop(input_image_path)
-    except HTTPException as e:
-        raise e
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Unexpected error: {str(e)}")
+# Set up the Gradio interface
+iface = gr.Interface(
+    fn=detect_and_crop,
+    inputs=gr.Image(type="pil"),  # Input type
+    outputs=gr.Image(type="numpy"),  # Output type
+    title="YOLOv7 Object Detection",
+    description="Upload an image for object detection and cropping."
+)
 
-    # Return the output image as a response
-    return FileResponse(output_image_path, media_type="image/jpeg")
+# Launch the app
+iface.launch()