Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import CLIPModel, CLIPProcessor | |
| from PIL import Image | |
| # Step 1: Load Fine-Tuned Model from Hugging Face Model Hub | |
| model_name = "quadranttechnologies/retail-content-safety-clip-finetuned" | |
| print("Initializing the application...") | |
| try: | |
| print("Loading the model from Hugging Face Model Hub...") | |
| model = CLIPModel.from_pretrained(model_name, trust_remote_code=True) | |
| processor = CLIPProcessor.from_pretrained(model_name) | |
| print("Model and processor loaded successfully.") | |
| except Exception as e: | |
| print(f"Error loading the model or processor: {e}") | |
| raise RuntimeError(f"Failed to load model: {e}") | |
| # Step 2: Define the Inference Function | |
| def classify_image(image): | |
| """ | |
| Classify an image as 'safe' or 'unsafe' and return probabilities. | |
| Args: | |
| image (PIL.Image.Image): Uploaded image. | |
| Returns: | |
| str: Predicted category ("safe" or "unsafe"). | |
| dict: Probabilities for "safe" and "unsafe". | |
| """ | |
| try: | |
| print("Starting image classification...") | |
| # Check if the image is valid | |
| if image is None: | |
| raise ValueError("No image provided. Please upload a valid image.") | |
| if not hasattr(image, "convert"): | |
| raise ValueError("Uploaded file is not a valid image format.") | |
| # Define main categories | |
| categories = ["safe", "unsafe"] | |
| print(f"Categories: {categories}") | |
| # Process the image | |
| print("Processing the image with the processor...") | |
| inputs = processor(text=categories, images=image, return_tensors="pt", padding=True) | |
| print(f"Processed inputs: {inputs}") | |
| # Perform inference | |
| print("Running model inference...") | |
| outputs = model(**inputs) | |
| print(f"Model outputs: {outputs}") | |
| # Calculate probabilities | |
| logits_per_image = outputs.logits_per_image # Image-text similarity scores | |
| probs = logits_per_image.softmax(dim=1) # Convert logits to probabilities | |
| print(f"Probabilities: {probs}") | |
| # Extract probabilities for each category | |
| safe_prob = probs[0][0].item() * 100 # Safe percentage | |
| unsafe_prob = probs[0][1].item() * 100 # Unsafe percentage | |
| # Determine the predicted category | |
| predicted_category = "safe" if safe_prob > unsafe_prob else "unsafe" | |
| print(f"Predicted category: {predicted_category}") | |
| # Return the predicted category and probabilities | |
| return predicted_category, {"safe": f"{safe_prob:.2f}%", "unsafe": f"{unsafe_prob:.2f}%"} | |
| except Exception as e: | |
| print(f"Error during classification: {e}") | |
| return f"Error: {str(e)}", {} | |
| # Step 3: Set Up Gradio Interface | |
| iface = gr.Interface( | |
| fn=classify_image, | |
| inputs=gr.Image(type="pil"), | |
| outputs=[ | |
| gr.Textbox(label="Predicted Category"), # Display the predicted category prominently | |
| gr.Label(label="Probabilities"), # Display probabilities with a progress bar | |
| ], | |
| title="Content Safety Classification", | |
| description="Upload an image to classify it as 'safe' or 'unsafe' with corresponding probabilities.", | |
| ) | |
| # Step 4: Launch Gradio Interface | |
| if __name__ == "__main__": | |
| print("Launching Gradio interface...") | |
| iface.launch() | |