musk12/Car-Segmentation-Mask
0
1# import os2# import torch3# import torchvision.transforms as T4# import torchvision.transforms.functional as TF5# import numpy as np6# from PIL import Image7# from flask import Flask, render_template, request, send_file, abort8 9# app = Flask(__name__)10 11# device = "cuda" if torch.cuda.is_available() else "cpu"12 13# # Load model (assuming UNet is defined in unet.py)14# def load_model():15# try:16# from unet import UNet17# model = UNet().to(device)18# model_path = "unet_car_final.pth"19# if not os.path.exists(model_path):20# raise FileNotFoundError(f"Model file {model_path} not found")21# model.load_state_dict(torch.load(model_path, map_location=device))22# model.eval()23# return model24# except Exception as e:25# print(f"Error loading model: {e}")26# raise27 28# try:29# model = load_model()30# except Exception as e:31# print(f"Model loading failed: {e}")32# model = None33 34# # Image transforms35# img_transform = T.Compose([36# T.Resize((256, 256)),37# T.ToTensor(),38# T.Normalize(mean=[0.485, 0.456, 0.406],39# std=[0.229, 0.224, 0.225])40# ])41 42# TMP_FOLDER = "/tmp"43# os.makedirs(TMP_FOLDER, exist_ok=True)44 45# # Route to serve files from /tmp46# @app.route('/tmp/<filename>')47# def serve_tmp_file(filename):48# file_path = os.path.join(TMP_FOLDER, filename)49# if os.path.exists(file_path):50# return send_file(file_path)51# else:52# print(f"File not found: {file_path}")53# abort(404)54 55# @app.route("/", methods=["GET", "POST"])56# def index():57# orig = None58# mask = None59# overlay = None60# error = None61 62# # Check for existing input image63# img_path = os.path.join(TMP_FOLDER, "input.jpg")64# if os.path.exists(img_path):65# orig = "/tmp/input.jpg"66# print(f"Found existing image: {img_path}")67 68# if request.method == "POST":69# # Handle image upload70# if "image" in request.files:71# file = request.files["image"]72# if file.filename == "":73# error = "No file selected"74# print(error)75# return render_template("index.html", error=error, orig=orig, mask=mask, overlay=overlay)76 77# try:78# # Save uploaded image to /tmp79# file.save(img_path)80# print(f"Image saved to: {img_path}")81# orig = "/tmp/input.jpg"82 83# # Clear previous results in /tmp84# for path in [os.path.join(TMP_FOLDER, "mask.png"), os.path.join(TMP_FOLDER, "overlay.png")]:85# if os.path.exists(path):86# os.remove(path)87# print(f"Removed: {path}")88# except Exception as e:89# error = f"Error uploading image: {str(e)}"90# print(f"Upload error: {e}")91# return render_template("index.html", error=error, orig=orig, mask=mask, overlay=overlay)92 93# # Handle segmentation94# if "segment" in request.form:95# if not os.path.exists(img_path):96# error = "No image available for segmentation"97# print(f"Segmentation error: Image not found at {img_path}")98# return render_template("index.html", error=error, orig=orig, mask=mask, overlay=overlay)99 100# try:101# if model is None:102# raise ValueError("Model not loaded")103 104# image = Image.open(img_path).convert("RGB")105# input_tensor = img_transform(image).unsqueeze(0).to(device)106 107# # Predict108# with torch.no_grad():109# output = model(input_tensor)110# pred_mask = torch.sigmoid(output)111# pred_mask = (pred_mask > 0.5).float()112 113# # Resize mask back to original image size114# mask_resized = TF.resize(115# TF.to_pil_image(pred_mask.squeeze().cpu()),116# size=image.size[::-1],117# interpolation=Image.NEAREST118# )119 120# # Save mask to /tmp121# mask_path = os.path.join(TMP_FOLDER, "mask.png")122# mask_resized.save(mask_path)123# print(f"Mask saved to: {mask_path}")124 125# # Create overlay126# mask_np = np.array(mask_resized)127# overlay = np.array(image).copy()128# overlay[mask_np > 128] = [255, 0, 0]129# overlay_img = Image.fromarray(overlay)130# overlay_path = os.path.join(TMP_FOLDER, "overlay.png")131# overlay_img.save(overlay_path)132# print(f"Overlay saved to: {overlay_path}")133 134# mask = "/tmp/mask.png"135# overlay = "/tmp/overlay.png"136# except Exception as e:137# error = f"Error during segmentation: {str(e)}"138# print(f"Segmentation error: {e}")139# return render_template("index.html", error=error, orig=orig, mask=mask, overlay=overlay)140 141# return render_template("index.html", orig=orig, mask=mask, overlay=overlay, error=error)142 143# if __name__ == "__main__":144# app.run(debug=True)145 146 147 148import os149import torch150import torchvision.transforms as T151import torchvision.transforms.functional as TF152import numpy as np153from PIL import Image154from flask import Flask, render_template, request, send_file, abort155 156app = Flask(__name__)157 158device = "cuda" if torch.cuda.is_available() else "cpu"159 160# Load model (assuming UNet is defined in unet.py)161def load_model():162 try:163 from unet import UNet164 model = UNet().to(device)165 model_path = "unet_car_final.pth"166 if not os.path.exists(model_path):167 raise FileNotFoundError(f"Model file {model_path} not found")168 model.load_state_dict(torch.load(model_path, map_location=device))169 model.eval()170 return model171 except Exception as e:172 print(f"Error loading model: {e}")173 raise174 175try:176 model = load_model()177except Exception as e:178 print(f"Model loading failed: {e}")179 model = None180 181# Image transforms182img_transform = T.Compose([183 T.Resize((256, 256)),184 T.ToTensor(),185 T.Normalize(mean=[0.485, 0.456, 0.406],186 std=[0.229, 0.224, 0.225])187])188 189TMP_FOLDER = "/tmp"190os.makedirs(TMP_FOLDER, exist_ok=True)191 192# Route to serve files from /tmp193@app.route('/tmp/<filename>')194def serve_tmp_file(filename):195 file_path = os.path.join(TMP_FOLDER, filename)196 if os.path.exists(file_path):197 return send_file(file_path)198 else:199 print(f"File not found: {file_path}")200 abort(404)201 202@app.route("/", methods=["GET", "POST"])203def index():204 orig = None205 mask = None206 overlay = None207 error = None208 209 if request.method == "GET":210 # Clear all relevant files in /tmp when a user accesses the root route211 for filename in ["input.jpg", "mask.png", "overlay.png"]:212 file_path = os.path.join(TMP_FOLDER, filename)213 if os.path.exists(file_path):214 try:215 os.remove(file_path)216 print(f"Cleared file: {file_path}")217 except Exception as e:218 print(f"Error clearing file {file_path}: {e}")219 220 # Check for existing input image (will be None since we cleared /tmp/input.jpg)221 img_path = os.path.join(TMP_FOLDER, "input.jpg")222 if os.path.exists(img_path):223 orig = "/tmp/input.jpg"224 print(f"Found existing image: {img_path}")225 226 if request.method == "POST":227 # Handle image upload228 if "image" in request.files:229 file = request.files["image"]230 if file.filename == "":231 error = "No file selected"232 print(error)233 return render_template("index.html", error=error, orig=orig, mask=mask, overlay=overlay)234 235 try:236 # Save uploaded image to /tmp237 file.save(img_path)238 print(f"Image saved to: {img_path}")239 orig = "/tmp/input.jpg"240 241 # Clear previous results in /tmp242 for path in [os.path.join(TMP_FOLDER, "mask.png"), os.path.join(TMP_FOLDER, "overlay.png")]:243 if os.path.exists(path):244 os.remove(path)245 print(f"Removed: {path}")246 except Exception as e:247 error = f"Error uploading image: {str(e)}"248 print(f"Upload error: {e}")249 return render_template("index.html", error=error, orig=orig, mask=mask, overlay=overlay)250 251 # Handle segmentation252 if "segment" in request.form:253 if not os.path.exists(img_path):254 error = "No image available for segmentation"255 print(f"Segmentation error: Image not found at {img_path}")256 return render_template("index.html", error=error, orig=orig, mask=mask, overlay=overlay)257 258 try:259 if model is None:260 raise ValueError("Model not loaded")261 262 image = Image.open(img_path).convert("RGB")263 input_tensor = img_transform(image).unsqueeze(0).to(device)264 265 # Predict266 with torch.no_grad():267 output = model(input_tensor)268 pred_mask = torch.sigmoid(output)269 pred_mask = (pred_mask > 0.5).float()270 271 # Resize mask back to original image size272 mask_resized = TF.resize(273 TF.to_pil_image(pred_mask.squeeze().cpu()),274 size=image.size[::-1],275 interpolation=Image.NEAREST276 )277 278 # Save mask to /tmp279 mask_path = os.path.join(TMP_FOLDER, "mask.png")280 mask_resized.save(mask_path)281 print(f"Mask saved to: {mask_path}")282 283 # Create overlay284 mask_np = np.array(mask_resized)285 overlay = np.array(image).copy()286 overlay[mask_np > 128] = [255, 0, 0]287 overlay_img = Image.fromarray(overlay)288 overlay_path = os.path.join(TMP_FOLDER, "overlay.png")289 overlay_img.save(overlay_path)290 print(f"Overlay saved to: {overlay_path}")291 292 mask = "/tmp/mask.png"293 overlay = "/tmp/overlay.png"294 except Exception as e:295 error = f"Error during segmentation: {str(e)}"296 print(f"Segmentation error: {e}")297 return render_template("index.html", error=error, orig=orig, mask=mask, overlay=overlay)298 299 return render_template("index.html", orig=orig, mask=mask, overlay=overlay, error=error)300 301if __name__ == "__main__":302 app.run(debug=True)