Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from ultralytics import YOLO | |
| import cv2 | |
| import numpy as np | |
| import tempfile | |
| import os | |
| # Load model | |
| print("Loading model...") | |
| model = YOLO('best.pt') | |
| print("Model loaded successfully") | |
| def process_image(image): | |
| try: | |
| print("\nProcessing image...") | |
| # Resize to match training size | |
| processed_image = cv2.resize(image, (1024, 1024)) | |
| # Run detection | |
| results = model(processed_image, conf=0.25) | |
| if results[0].obb is not None and len(results[0].obb) > 0: | |
| obb_results = results[0].obb | |
| # Count detections by class | |
| caries_count = 0 | |
| non_caries_count = 0 | |
| # Process detections | |
| for i in range(len(obb_results)): | |
| cls = int(obb_results.cls[i]) | |
| if cls == 0: # assuming 0 is caries | |
| caries_count += 1 | |
| else: | |
| non_caries_count += 1 | |
| # Get annotated image | |
| annotated_image = results[0].plot( | |
| conf=True, | |
| line_width=2, | |
| font_size=15, | |
| labels=True | |
| ) | |
| # Resize back to original size | |
| annotated_image = cv2.resize(annotated_image, (image.shape[1], image.shape[0])) | |
| summary = f"Found {caries_count} cavities and {non_caries_count} normal teeth" | |
| else: | |
| annotated_image = image | |
| summary = "No detections found" | |
| return annotated_image, summary | |
| except Exception as e: | |
| print(f"Error in image processing: {str(e)}") | |
| return image, f"Error: {str(e)}" | |
| def process_video(video_path): | |
| try: | |
| print("Processing video...") | |
| cap = cv2.VideoCapture(video_path) | |
| # Get video properties | |
| width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
| total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| # Calculate frame interval (7 seconds * fps) | |
| interval = 7 * fps | |
| # Create temporary output file | |
| temp_dir = tempfile.gettempdir() | |
| output_path = os.path.join(temp_dir, 'output_video.mp4') | |
| # Initialize video writer | |
| output = cv2.VideoWriter( | |
| output_path, | |
| cv2.VideoWriter_fourcc(*'mp4v'), | |
| fps, | |
| (width, height) | |
| ) | |
| frame_count = 0 | |
| last_processed_frame = None | |
| frames_since_last_process = 0 | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Process frame if it's at the interval or first frame | |
| if frame_count == 0 or frames_since_last_process >= interval: | |
| print(f"Processing frame at {frame_count / fps:.1f} seconds") | |
| # Process frame | |
| processed_frame = cv2.resize(frame, (1024, 1024)) | |
| results = model(processed_frame, conf=0.25) | |
| if results[0].obb is not None: | |
| # Draw detections | |
| annotated_frame = results[0].plot( | |
| conf=True, | |
| line_width=2, | |
| font_size=15, | |
| labels=True | |
| ) | |
| annotated_frame = cv2.resize(annotated_frame, (width, height)) | |
| last_processed_frame = annotated_frame | |
| else: | |
| last_processed_frame = frame | |
| frames_since_last_process = 0 | |
| # Write the last processed frame | |
| if last_processed_frame is not None: | |
| output.write(last_processed_frame) | |
| else: | |
| output.write(frame) | |
| frame_count += 1 | |
| frames_since_last_process += 1 | |
| print(f"Processed {frame_count}/{total_frames} frames", end='\r') | |
| cap.release() | |
| output.release() | |
| summary = f"Processed video with {total_frames} frames\nDetected frames at 7-second intervals" | |
| return output_path, summary | |
| except Exception as e: | |
| print(f"Error in video processing: {str(e)}") | |
| return None, f"Error: {str(e)}" | |
| # Create Gradio interface with examples | |
| with gr.Blocks(title="Dental Cavity Detection") as demo: | |
| gr.Markdown(""" | |
| # Dental Cavity Detection System | |
| Upload or select an example to detect dental cavities | |
| """) | |
| with gr.Tabs(): | |
| with gr.Tab("Image Detection"): | |
| with gr.Row(): | |
| # Input Column | |
| with gr.Column(): | |
| image_input = gr.Image(label="Upload Image") | |
| image_button = gr.Button("Detect Cavities", variant="primary") | |
| # Output Column | |
| with gr.Column(): | |
| image_output = gr.Image(label="Detection Result") | |
| image_summary = gr.Textbox(label="Detection Summary") | |
| # Add examples using Gradio's examples feature | |
| gr.Examples( | |
| examples=[ | |
| ["image1.jpg"], | |
| ["image2.jpg"] | |
| ], | |
| inputs=image_input, | |
| label="Example Images - Click to use" | |
| ) | |
| # Process button click | |
| image_button.click( | |
| fn=process_image, | |
| inputs=image_input, | |
| outputs=[image_output, image_summary] | |
| ) | |
| with gr.Tab("Video Detection"): | |
| with gr.Row(): | |
| # Input Column | |
| with gr.Column(): | |
| video_input = gr.Video(label="Upload Video") | |
| video_button = gr.Button("Process Video", variant="primary") | |
| # Output Column | |
| with gr.Column(): | |
| video_output = gr.Video(label="Processed Video") | |
| gr.Markdown("""⚠️ Note: Video preview may not work in browser. | |
| Please use the download button below the video to view results.""") | |
| video_summary = gr.Textbox(label="Processing Summary") | |
| # Add video example | |
| gr.Examples( | |
| examples=[ | |
| ["video1.mp4"] | |
| ], | |
| inputs=video_input, | |
| label="Example Video - Click to use" | |
| ) | |
| # Process button click | |
| video_button.click( | |
| fn=process_video, | |
| inputs=video_input, | |
| outputs=[video_output, video_summary] | |
| ) | |
| # Add footer with instructions | |
| gr.Markdown(""" | |
| ### Instructions: | |
| 1. Choose Image or Video tab | |
| 2. Upload your own file or click an example below | |
| 3. Click the detection button to process | |
| """) | |
| if __name__ == "__main__": | |
| demo.launch(debug=True) |