Upload 2 files

app.py

@@ -3,52 +3,63 @@ import requests
 import numpy as np
 from PIL import Image
 import io
+from ultralytics import YOLO
+import cv2  
 
-# Function to call the Hugging Face model
-def call_model(image):
-    # TODO: Replace with your actual Hugging Face model endpoint
-    API_URL = "YOUR_HUGGINGFACE_MODEL_ENDPOINT"
-    headers = {"Authorization": "Bearer YOUR_API_TOKEN"}
-    
-    # Convert image to bytes
-    img_byte_arr = io.BytesIO()
-    image.save(img_byte_arr, format='PNG')
-    img_byte_arr = img_byte_arr.getvalue()
-    
-    # Make API call
-    response = requests.post(API_URL, headers=headers, data=img_byte_arr)
-    
-    # Process response and modify image
-    # TODO: Implement your image modification logic based on model results
-    # This is a placeholder that just returns the original image
-    return image
+# Load the YOLO model at startup
+try:
+    model = YOLO('modelo_epoch_50.pt')
+    print("Model loaded successfully")
+except Exception as e:
+    print(f"Error loading model: {str(e)}")
+    model = None
 
 # Function to process the captured image
-def process_image(image):
+def process_image(image):    
+    # Check if image is None
     if image is None:
-        return None
+        raise gr.Error("Please take a picture first before analyzing!")
     
-    # Convert to PIL Image if it's not already
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
+    # Convert image to RGB if it's not already
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
     
-    # Call the model and get modified image
-    modified_image = call_model(image)
+    # Run inference
+    results = model.predict(image, save=True, conf=0.5)
+
+    # Print the results
+    print("Model predictions:", results[0].boxes)
+
+    # Convert PIL Image to numpy array for OpenCV
+    image_cv = np.array(image)
+    # Convert RGB to BGR (OpenCV uses BGR format)
+    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
+
+    # Draw bounding boxes of the prediction
+    boxes = results[0].boxes
+    for box in boxes:
+        b = box.xyxy[0]  # Bounding box coordinates
+        c = box.cls  # Predicted class
+        confidence = box.conf
+
+        x1, y1, x2, y2 = map(int, b)
+        cv2.rectangle(image_cv, (x1, y1), (x2, y2), (255, 0, 0), 2)  # Blue for prediction
+
+        label = f"{results[0].names[int(c)]} {confidence.item():.2f}"
+        cv2.putText(image_cv, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)
+        print(label)
+
+    # Convert back to RGB for display if needed
+    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
+    # Convert back to PIL Image
+    final_image = Image.fromarray(image_cv)
+    # Save the final image to current directory, overwriting if exists
+    final_image.save('FINAL.jpg', 'JPEG', quality=95)
     
-    return modified_image
+    return final_image
 
 # Create the Gradio interface
 with gr.Blocks() as demo:
-    # Title and main description
-    gr.Markdown("# Target Analyzer")
-    gr.Markdown("### Instructions:")
-    gr.Markdown("""
-    1. Click the '📸 Take Photo' button in the camera view to capture an image
-    2. Click 'Analyze' to process the image
-    3. View the results in the 'Analyzed Image' section
-    4. Use 'Reset' to start over
-    """)
-    
     with gr.Row():
         # Left column for camera and controls
         with gr.Column(scale=1):
@@ -80,4 +91,4 @@ with gr.Blocks() as demo:
 
 # Launch the interface
 if __name__ == "__main__":
-    demo.launch() 
+    demo.launch(share=True) 

requirements.txt

@@ -1,4 +1,5 @@
 gradio>=4.0.0
 Pillow>=10.0.0
 requests>=2.31.0
-numpy>=1.24.0 
+numpy>=1.24.0
+ultralytics>=8.0.0