Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import cv2 | |
| import numpy as np | |
| from ultralytics import YOLO | |
| from PIL import Image | |
| # Load YOLO models | |
| try: | |
| front_model = YOLO("front_cnic_model.pt") | |
| back_model = YOLO("back_cnic_model.pt") | |
| except Exception as e: | |
| print(f"Error loading models: {e}") | |
| def preprocess_image(image): | |
| """Convert PIL Image to OpenCV format.""" | |
| try: | |
| img = np.array(image) | |
| img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) | |
| return img | |
| except Exception as e: | |
| return None, f"Error preprocessing image: {e}" | |
| def extract_regions(image, boxes, model_names): | |
| """Extract regions of interest (ROIs) by label from YOLO detections.""" | |
| results = {} | |
| for box in boxes: | |
| try: | |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) | |
| cls_name = model_names[int(box.cls)] | |
| roi = image[y1:y2, x1:x2] | |
| results[cls_name] = f"Detected region at [{x1}, {y1}, {x2}, {y2}]" | |
| except Exception as e: | |
| results[cls_name] = f"Region extraction error: {e}" | |
| return results | |
| def process_front_cnic(image): | |
| """Process front CNIC image to extract fields.""" | |
| if image is None: | |
| return {"error": "No image provided"} | |
| img = preprocess_image(image) | |
| if isinstance(img, tuple): | |
| return {"error": img[1]} | |
| try: | |
| results = front_model.predict(img, conf=0.5) | |
| extracted_info = extract_regions(img, results[0].boxes, front_model.names) | |
| return extracted_info if extracted_info else {"error": "No objects detected"} | |
| except Exception as e: | |
| return {"error": f"Front CNIC processing error: {e}"} | |
| def process_back_cnic(image): | |
| """Process back CNIC image to extract ROIs for Barcode and CNIC number.""" | |
| if image is None: | |
| return {"error": "No image provided"} | |
| img = preprocess_image(image) | |
| if isinstance(img, tuple): | |
| return {"error": img[1]} | |
| try: | |
| results = back_model.predict(img, conf=0.5) | |
| extracted_info = extract_regions(img, results[0].boxes, back_model.names) | |
| return extracted_info if extracted_info else {"error": "No objects detected"} | |
| except Exception as e: | |
| return {"error": f"Back CNIC processing error: {e}"} | |
| # Gradio Interface | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown( | |
| """ | |
| # CNIC Detection (YOLO-only) | |
| Upload front and back CNIC images to detect regions of interest. | |
| - **Front CNIC**: Detects regions such as Name, DOB, Father's Name, etc. | |
| - **Back CNIC**: Detects Barcode and CNIC number regions (no OCR). | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("## Front CNIC") | |
| front_input = gr.Image(type="pil", label="Upload Front CNIC Image") | |
| front_button = gr.Button("Process Front CNIC") | |
| front_output = gr.JSON(label="Detected Regions") | |
| with gr.Column(): | |
| gr.Markdown("## Back CNIC") | |
| back_input = gr.Image(type="pil", label="Upload Back CNIC Image") | |
| back_button = gr.Button("Process Back CNIC") | |
| back_output = gr.JSON(label="Detected Regions") | |
| front_button.click( | |
| fn=process_front_cnic, | |
| inputs=front_input, | |
| outputs=front_output | |
| ) | |
| back_button.click( | |
| fn=process_back_cnic, | |
| inputs=back_input, | |
| outputs=back_output | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |