Update app.py

app.py

@@ -1,69 +1,26 @@
-import gradio as gr
-import cv2
-import torch
-from transformers import DetrImageProcessor, DetrForObjectDetection
+import numpy as np 
 from PIL import Image
-import tempfile
-import os
-
-# Load the DETR processor and model
-processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
-model = DetrForObjectDetection.from_pretrained("ArrayDice/Vehicle_Detection_Model_Zoom")
-
-# Function to process video and detect vehicles
-def detect_vehicles_in_video(video_path):
-    # Load video
-    cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-    # Temporary output video file
-    temp_video = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-    out = cv2.VideoWriter(temp_video.name, cv2.VideoWriter_fourcc(*"mp4v"), fps, (width, height))
-
-    # Process each frame
-    while cap.isOpened():
-        ret, frame = cap.read()
-        if not ret:
-            break
-
-        # Convert frame to PIL for model processing
-        pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-
-        # Run object detection
-        inputs = processor(images=pil_image, return_tensors="pt")
-        outputs = model(**inputs)
-
-        # Set confidence threshold and draw boxes
-        for score, label, box in zip(outputs.logits.softmax(-1)[0], outputs.labels[0], outputs.boxes[0]):
-            if score.max() >= 0.9:  # Adjust confidence threshold as needed
-                x_min, y_min, x_max, y_max = map(int, box)
-                cv2.rectangle(frame, (x_min, y_min), (x_max, y_max), (0, 255, 0), 2)
-                label_text = f"Vehicle ({score.max():.2f})"
-                cv2.putText(frame, label_text, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
-
-        # Write annotated frame to output video
-        out.write(frame)
-
-    # Release resources
-    cap.release()
-    out.release()
-
-    return temp_video.name  # Return path to the annotated video
-
-# Gradio interface
-def vehicle_detection_gradio(video):
-    annotated_video_path = detect_vehicles_in_video(video.name)
-    return annotated_video_path
-
-# Set up Gradio interface
-iface = gr.Interface(
-    fn=vehicle_detection_gradio,
-    inputs=gr.Video(type="file", label="Upload a video for vehicle detection"),
-    outputs=gr.Video(label="Annotated Video with Detected Vehicles"),
-    title="Vehicle Detection Model",
-    description="Upload a video, and the model will detect and annotate vehicles in each frame."
-)
-
-iface.launch()
+from transformers import AutoImageProcessor, AutoModelForImageClassification
+
+# Load the model and image processor
+processor = AutoImageProcessor.from_pretrained("beingamit99/car_damage_detection")
+model = AutoModelForImageClassification.from_pretrained("beingamit99/car_damage_detection")
+
+# Load and process the image
+image = Image.open(IMAGE)
+inputs = processor(images=image, return_tensors="pt")
+
+# Make predictions
+outputs = model(**inputs)
+logits = outputs.logits.detach().cpu().numpy()
+predicted_class_id = np.argmax(logits)
+predicted_proba = np.max(logits)
+label_map = model.config.id2label
+predicted_class_name = label_map[predicted_class_id]
+
+# Print the results
+print(f"Predicted class: {predicted_class_name} (probability: {predicted_proba:.4f}")
+from transformers import pipeline
+#Create a classification pipeline
+pipe = pipeline("image-classification", model="beingamit99/car_damage_detection")
+pipe(IMAGE)