Create app.py

app.py

@@ -0,0 +1,83 @@
+import torch
+from transformers import AutoImageProcessor, AutoModelForObjectDetection
+from PIL import Image
+import cv2
+import numpy as np
+import time
+import gradio as gr
+
+# Device setup (GPU or CPU)
+device = 'cpu'
+if torch.cuda.is_available():
+    device = torch.device('cuda')
+elif torch.backends.mps.is_available():
+    device = torch.device('mps')
+
+# Load pre-trained model and image processor from Hugging Face
+ckpt = 'yainage90/fashion-object-detection'
+image_processor = AutoImageProcessor.from_pretrained(ckpt)
+model = AutoModelForObjectDetection.from_pretrained(ckpt).to(device)
+
+def detect_objects(frame):
+    """Detect objects in the video frame."""
+    # Convert the frame to PIL image
+    image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+
+    # Prepare inputs for the model
+    with torch.no_grad():
+        inputs = image_processor(images=[image], return_tensors="pt")
+        outputs = model(**inputs.to(device))
+        target_sizes = torch.tensor([[image.size[1], image.size[0]]])
+        results = image_processor.post_process_object_detection(outputs, threshold=0.4, target_sizes=target_sizes)[0]
+
+    # Extract the detected items
+    items = []
+    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        score = score.item()
+        label = label.item()
+        box = [i.item() for i in box]
+        print(f"{model.config.id2label[label]}: {round(score, 3)} at {box}")
+        items.append((score, label, box))
+
+    return items
+
+def process_image(image):
+    """Process the image uploaded via Gradio and return the result."""
+    # Convert the image to numpy array
+    frame = np.array(image)
+
+    # Detect objects (e.g., helmets) in the frame
+    items = detect_objects(frame)
+
+    if items:  # If objects are detected, save the data
+        save_data(frame, items)
+
+    return {"items_detected": items}
+
+def save_data(frame, items):
+    """Save image and extract plate number."""
+    filename = f"helmet_violation_{int(time.time())}.jpg"
+    cv2.imwrite(filename, frame)
+    
+    # Here, you'd extract plate numbers or process further
+    plate_number = extract_plate_number(frame)
+    save_to_database(filename, plate_number, items)
+
+def extract_plate_number(frame):
+    """Extract license plate number (simplified)."""
+    plate_number = "XYZ 1234"  # Replace with an actual license plate recognition method
+    return plate_number
+
+def save_to_database(image_filename, plate_number, items):
+    """Save the data (for simplicity, we just print it here)."""
+    print(f"Plate Number: {plate_number}, Image saved as {image_filename}")
+    print("Detected items:", items)
+
+# Define the Gradio interface
+interface = gr.Interface(fn=process_image,
+                         inputs=gr.inputs.Image(type="pil"),
+                         outputs=gr.outputs.JSON(),
+                         live=True)
+
+# Launch the Gradio app
+interface.launch(debug=True)