Spaces:
Sleeping
Sleeping
| import random # Import random for selecting examples | |
| import os # For reading environment variables | |
| import shutil # For directory cleanup | |
| import zipfile # For extracting model archives | |
| import pathlib # For path manipulations | |
| import tempfile # For creating temporary files/directories | |
| import numpy as np # For image processing | |
| # import pickle # No longer needed for loading from zip | |
| import gradio # For interactive UI | |
| import pandas # For tabular data handling | |
| import PIL.Image # For image I/O | |
| import huggingface_hub # For downloading model assets | |
| import autogluon.multimodal # For loading AutoGluon image classifier | |
| # Hardcoded Hub model (native zip) | |
| MODEL_REPO_ID = "jennifee/nnl_automl_model" # Updated model ID | |
| ZIP_FILENAME = "autogluon_image_predictor_dir.zip" # Updated filename | |
| HF_TOKEN = os.getenv("HF_TOKEN", None) | |
| # Local cache/extract dirs | |
| CACHE_DIR = pathlib.Path("hf_assets") | |
| EXTRACT_DIR = CACHE_DIR / "predictor_native" # Keep extract dir name for consistency | |
| # Download & load the native predictor | |
| def _prepare_predictor_dir() -> str: # Reverted function name | |
| # Clear the Hugging Face Hub cache to avoid caching issues - Keep for now, can remove if issues persist | |
| from huggingface_hub import delete_repo | |
| try: | |
| # Use the current MODEL_REPO_ID for deletion | |
| delete_repo(MODEL_REPO_ID, repo_type="model", token=HF_TOKEN) | |
| except Exception as e: | |
| print(f"Could not delete repo from cache (may not exist or unauthorized): {e}") | |
| CACHE_DIR.mkdir(parents=True, exist_ok=True) | |
| local_zip = huggingface_hub.hf_hub_download( | |
| repo_id=MODEL_REPO_ID, | |
| filename=ZIP_FILENAME, | |
| repo_type="model", | |
| token=HF_TOKEN, | |
| local_dir=str(CACHE_DIR), | |
| local_dir_use_symlinks=False, | |
| ) | |
| if EXTRACT_DIR.exists(): | |
| shutil.rmtree(EXTRACT_DIR) | |
| EXTRACT_DIR.mkdir(parents=True, exist_ok=True) | |
| with zipfile.ZipFile(local_zip, "r") as zf: | |
| zf.extractall(str(EXTRACT_DIR)) | |
| contents = list(EXTRACT_DIR.iterdir()) | |
| predictor_root = contents[0] if (len(contents) == 1 and contents[0].is_dir()) else EXTRACT_DIR | |
| return str(predictor_root) | |
| try: | |
| PREDICTOR_DIR = _prepare_predictor_dir() # Call the function to prepare the directory | |
| # PREDICTOR_FILE = _prepare_predictor_file() # Old call | |
| # Load the predictor from the extracted directory | |
| PREDICTOR = autogluon.multimodal.MultiModalPredictor.load(PREDICTOR_DIR) # Updated load | |
| # with open(PREDICTOR_FILE, 'rb') as f: # Old pickle load | |
| # PREDICTOR = pickle.load(f) # Old pickle load | |
| except Exception as e: | |
| print(f"Error loading predictor: {e}") | |
| PREDICTOR = None # Set predictor to None if loading fails | |
| # Explicit class labels (edit copy as desired) - Updated based on user input | |
| CLASS_LABELS = { | |
| 0: "👑 Face", | |
| 1: "🔢 Value", | |
| # Removed suit labels and Joker as per new mapping | |
| # 5: "Unknown" # Keep Unknown as a placeholder if needed | |
| } | |
| # Helper to map model class -> human label | |
| def _human_label(c): | |
| try: | |
| # Attempt to convert to integer first, then use get for safety | |
| ci = int(c) | |
| return CLASS_LABELS.get(ci, str(c)) | |
| except (ValueError, TypeError): | |
| # If conversion fails, try getting directly by key | |
| return CLASS_LABELS.get(c, str(c)) | |
| # Function to preprocess image and return processed image - Keep as is for now | |
| def preprocess_image_for_display(pil_img: PIL.Image.Image): | |
| if pil_img is None: | |
| return None | |
| # AutoGluon preprocessing (simplified, actual preprocessing is done internally by the predictor) | |
| # Here we resize for display purposes to show a consistent "processed" image | |
| processed_img = pil_img.resize((224, 224)) # Example size, adjust as needed | |
| return processed_img | |
| # Do the prediction! - Adjusting outputs to match the new model's likely output | |
| def do_predict(pil_img: PIL.Image.Image): | |
| # Make sure there's actually an image to work with and predictor is loaded | |
| if pil_img is None: | |
| # Returning None for the processed image output when input is None | |
| # Adjusting return values to match expected outputs: status, probabilities, processed image | |
| return "No image provided.", {}, None | |
| if PREDICTOR is None: | |
| # Returning None for the processed image output when predictor is not loaded | |
| # Adjusting return values to match expected outputs: status, probabilities, processed image | |
| return "Predictor not loaded. Please check the logs for errors.", {}, None | |
| # Basic validation (file type is handled by Gradio, checking size) | |
| # This is a placeholder; real size checks would be on file upload before PIL | |
| # For now, we'll just check if the image object is valid | |
| try: | |
| pil_img.verify() | |
| except Exception: | |
| # Returning None for the processed image output for invalid image | |
| # Adjusting return values to match expected outputs: status, probabilities, processed image | |
| return "Invalid image file.", {}, None | |
| # IF we have something to work with, save it and prepare the input | |
| tmpdir = pathlib.Path(tempfile.mkdtemp()) | |
| img_path = tmpdir / "input.png" | |
| pil_img.save(img_path) | |
| df = pandas.DataFrame({"image": [str(img_path)]}) # For AutoGluon expected input format | |
| # For class probabilities | |
| # Assuming predict_proba returns a DataFrame where columns are class labels | |
| proba_output = PREDICTOR.predict_proba(df) | |
| print(f"Type of proba_output: {type(proba_output)}") | |
| print(f"Content of proba_output: {proba_output}") | |
| # Assuming proba_output is a pandas DataFrame with class probabilities | |
| if not proba_output.empty: | |
| # Get probabilities for the first (and likely only) row | |
| proba_series = proba_output.iloc[0] | |
| # Convert to dictionary, mapping original labels to probabilities | |
| proba_dict = proba_series.to_dict() | |
| # For user-friendly column names | |
| # Map the original labels (keys in proba_dict) to human-friendly labels | |
| pretty_dict = { | |
| _human_label(k): float(v) for k, v in proba_dict.items() | |
| } | |
| else: | |
| # Handle case where predict_proba returns empty | |
| pretty_dict = {} | |
| # Generate processed image for display | |
| processed_img_display = preprocess_image_for_display(pil_img) | |
| # Return prediction result, probabilities, and the processed image for display | |
| # Ensure the number of return values matches the outputs in the Gradio interface | |
| return "Prediction Complete", pretty_dict, processed_img_display | |
| # Representative example images! These can be local or links. | |
| # Using the provided local file paths as examples | |
| EXAMPLES = [ | |
| ["./examples/WhatsApp Image 2025-09-12 at 22.05.40 (2).jpeg"], | |
| ["./examples/WhatsApp Image 2025-09-12 at 22.05.40 (3).jpeg"], | |
| ["./examples/WhatsApp Image 2025-09-12 at 22.05.40 (5).jpeg"] | |
| ] | |
| # Gradio UI | |
| with gradio.Blocks() as demo: | |
| # Provide an introduction - Updated for Playing Cards | |
| gradio.Markdown("# Playing Card Detection") | |
| gradio.Markdown(""" | |
| This is a simple app that demonstrates how to use an autogluon multimodal | |
| predictor in a gradio space to predict the type of playing card in a picture. To use, | |
| just upload a photo using the options below. The original and preprocessed | |
| images will be displayed, and the prediction results will appear automatically. | |
| """) | |
| with gradio.Row(): | |
| # Interface for the incoming image | |
| image_in = gradio.Image(type="pil", label="Input image", sources=["upload", "webcam", "clipboard"]) | |
| # Display preprocessed image | |
| image_processed_out = gradio.Image(type="pil", label="Preprocessed image") | |
| # Interface elements to show the result and probabilities | |
| # Adjusting num_top_classes if the model has more classes | |
| proba_pretty = gradio.Label(num_top_classes=len(CLASS_LABELS), label="Class probabilities") | |
| prediction_status = gradio.Textbox(label="Prediction Status") # Added Textbox for status | |
| # Expose key inference parameters (placeholder) | |
| # gradio.Markdown("## Inference Parameters") | |
| # with gradio.Row(): | |
| # Add parameters here if needed, e.e., a confidence threshold slider | |
| # confidence_threshold = gradio.Slider(minimum=0, maximum=1, value=0.5, label="Confidence Threshold") | |
| # Whenever a new image is uploaded, trigger the prediction directly | |
| # Wrap do_predict in a lambda to ensure only the image input is passed | |
| # Ensure outputs match the return values of do_predict | |
| image_in.change( | |
| fn=lambda img: do_predict(img), | |
| inputs=[image_in], | |
| outputs=[prediction_status, proba_pretty, image_processed_out] | |
| ) | |
| # For clickable example images - ADDED BACK | |
| if EXAMPLES: # Only show examples if any were successfully fetched | |
| gradio.Examples( | |
| examples=EXAMPLES, | |
| inputs=[image_in], | |
| label="Representative examples", | |
| examples_per_page=8, | |
| cache_examples=False, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |