Update app.py

app.py

@@ -1,92 +1,43 @@
-from fastapi import FastAPI, File, UploadFile
-import uvicorn
-from typing import List
-from io import BytesIO
-import numpy as np
-import rasterio
-from pydantic import BaseModel
-import torch
+######### pull files
+import os
 from huggingface_hub import hf_hub_download
+config_path=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-multi-temporal-crop-classification", 
+                            filename="multi_temporal_crop_classification_Prithvi_100M.py", 
+                            token=os.environ.get("token"))
+ckpt=hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-multi-temporal-crop-classification", 
+                     filename='multi_temporal_crop_classification_Prithvi_100M.pth', 
+                     token=os.environ.get("token"))
+##########
+import argparse
 from mmcv import Config
-from mmseg.apis import init_segmentor
-import gradio as gr
-from functools import partial
-import time
-import os
 
-# Initialize the FastAPI app
-app = FastAPI()
+from mmseg.models import build_segmentor
 
-# Load the model and config
-config_path = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-multi-temporal-crop-classification", 
-                              filename="multi_temporal_crop_classification_Prithvi_100M.py", 
-                              token=os.environ.get("token"))
-ckpt = hf_hub_download(repo_id="ibm-nasa-geospatial/Prithvi-100M-multi-temporal-crop-classification", 
-                       filename='multi_temporal_crop_classification_Prithvi_100M.pth', 
-                       token=os.environ.get("token"))
+from mmseg.datasets.pipelines import Compose, LoadImageFromFile
 
-config = Config.fromfile(config_path)
-config.model.backbone.pretrained = None
-model = init_segmentor(config, ckpt, device='cpu')
+import rasterio
+import torch
 
-# Use the test pipeline directly
-custom_test_pipeline = model.cfg.data.test.pipeline
+from mmseg.apis import init_segmentor
 
-# Define the input/output model for FastAPI
-class PredictionOutput(BaseModel):
-    t1: List[float]
-    t2: List[float]
-    t3: List[float]
-    prediction: List[float]
+from mmcv.parallel import collate, scatter
 
-# Define the inference function
-def inference_on_file(file_path, model, custom_test_pipeline):
-    with rasterio.open(file_path) as src:
-        img = src.read()
+import numpy as np
+import glob
+import os
 
-    # Apply preprocessing using the custom pipeline
-    processed_img = apply_pipeline(custom_test_pipeline, img)
-    
-    # Run inference
-    output = model.inference(processed_img)
-    
-    # Post-process the output to get the RGB and prediction images
-    rgb1 = postprocess_output(output[0])
-    rgb2 = postprocess_output(output[1])
-    rgb3 = postprocess_output(output[2])
-    
-    return rgb1, rgb2, rgb3, output
+import time
 
-def apply_pipeline(pipeline, img):
-    # Implement your custom pipeline processing here
-    # This could include normalization, resizing, etc.
-    return img
+import numpy as np
+import gradio as gr
+from functools import partial
 
-def postprocess_output(output):
-    # Convert the model's output into an RGB image or other formats as needed
-    return output
+import pdb
 
-@app.post("/predict/", response_model=PredictionOutput)
-async def predict(file: UploadFile = File(...)):
-    # Read the uploaded file
-    target_image = BytesIO(await file.read())
-    
-    # Save the file temporarily if needed
-    with open("temp_image.tif", "wb") as f:
-        f.write(target_image.getvalue())
+import matplotlib.pyplot as plt
+
+from skimage import exposure 
 
-    # Run the prediction
-    rgb1, rgb2, rgb3, output = inference_on_file("temp_image.tif", model, custom_test_pipeline)
-    
-    # Return the results
-    return {
-        "t1": rgb1.tolist(),
-        "t2": rgb2.tolist(),
-        "t3": rgb3.tolist(),
-        "prediction": output.tolist()
-    }
-
-# Optional: Serve the Gradio interface (if you still want to use it with FastAPI)
 cdl_color_map = [{'value': 1, 'label': 'Natural vegetation', 'rgb': (233,255,190)},
                  {'value': 2, 'label': 'Forest', 'rgb': (149,206,147)},
                  {'value': 3, 'label': 'Corn', 'rgb': (255,212,0)},
@@ -316,8 +267,4 @@ with gr.Blocks() as demo:
             gr.Image(value='Legend.png', image_mode='RGB', show_label=False)
     
 
-demo.launch() 
-
-if __name__ == "__main__":
-    run_gradio_interface()
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+demo.launch()