Spaces:
Sleeping
Sleeping
| import os | |
| import cv2 | |
| import numpy as np | |
| from cv2 import dnn_superres | |
| class AIStudioService: | |
| """ | |
| Handles Deep Learning tasks: | |
| - Colorization (Caffe) | |
| - Super Resolution (EDSR x4) | |
| - Restoration (Denoising) | |
| """ | |
| def __init__(self): | |
| self.models_dir = 'models' | |
| # Colorization Paths | |
| self.proto_path = os.path.join(self.models_dir, 'colorization_deploy_v2.prototxt') | |
| self.model_path = os.path.join(self.models_dir, 'colorization_release_v2.caffemodel') | |
| self.points_path = os.path.join(self.models_dir, 'pts_in_hull.npy') | |
| # Super Resolution Path | |
| self.edsr_path = os.path.join(self.models_dir, 'EDSR_x4.pb') | |
| def process_request(self, image_path, output_path, task): | |
| try: | |
| if not os.path.exists(image_path): | |
| raise FileNotFoundError("Input file not found") | |
| if task == 'colorize': | |
| return self._colorize_image(image_path, output_path) | |
| elif task == 'upscale': | |
| return self._upscale_image(image_path, output_path) | |
| elif task == 'restore': | |
| return self._restore_image(image_path, output_path) | |
| else: | |
| raise ValueError(f"Unknown AI task: {task}") | |
| except Exception as e: | |
| print(f"AI Service Error: {e}") | |
| return False, str(e) | |
| def _colorize_image(self, input_path, output_path): | |
| if not os.path.exists(self.model_path): | |
| raise FileNotFoundError("Colorization models missing") | |
| net = cv2.dnn.readNetFromCaffe(self.proto_path, self.model_path) | |
| pts = np.load(self.points_path) | |
| class8 = net.getLayerId("class8_ab") | |
| conv8 = net.getLayerId("conv8_313_rh") | |
| pts = pts.transpose().reshape(2, 313, 1, 1) | |
| net.getLayer(class8).blobs = [pts.astype("float32")] | |
| net.getLayer(conv8).blobs = [np.full([1, 313], 2.606, dtype="float32")] | |
| img = cv2.imread(input_path) | |
| normalized = img.astype("float32") / 255.0 | |
| lab = cv2.cvtColor(normalized, cv2.COLOR_BGR2LAB) | |
| resized = cv2.resize(lab, (224, 224)) | |
| L = cv2.split(resized)[0] | |
| L -= 50 | |
| net.setInput(cv2.dnn.blobFromImage(L)) | |
| ab = net.forward()[0, :, :, :].transpose((1, 2, 0)) | |
| ab = cv2.resize(ab, (img.shape[1], img.shape[0])) | |
| L_orig = cv2.split(lab)[0] | |
| colorized = np.concatenate((L_orig[:, :, np.newaxis], ab), axis=2) | |
| colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2BGR) | |
| colorized = np.clip(colorized, 0, 1) | |
| colorized = (255 * colorized).astype("uint8") | |
| cv2.imwrite(output_path, colorized) | |
| return True, output_path | |
| def _upscale_image(self, input_path, output_path): | |
| if not os.path.exists(self.edsr_path): | |
| raise FileNotFoundError("EDSR_x4.pb model missing") | |
| img = cv2.imread(input_path) | |
| h, w = img.shape[:2] | |
| # Safety Check: Resize if too big before upscaling | |
| max_dim = 1024 | |
| if max(h, w) > max_dim: | |
| scale = max_dim / max(h, w) | |
| new_w, new_h = int(w * scale), int(h * scale) | |
| img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA) | |
| sr = dnn_superres.DnnSuperResImpl_create() | |
| sr.readModel(self.edsr_path) | |
| sr.setModel("edsr", 4) | |
| result = sr.upsample(img) | |
| cv2.imwrite(output_path, result) | |
| return True, output_path | |
| def _restore_image(self, input_path, output_path): | |
| """ | |
| Removes noise from image (Denoising). | |
| """ | |
| img = cv2.imread(input_path) | |
| # h = strength of filter. Big h = perfectly smooth but blurred details. | |
| # 10 is a good balance. | |
| result = cv2.fastNlMeansDenoisingColored(img, None, 10, 10, 7, 21) | |
| cv2.imwrite(output_path, result) | |
| return True, output_path | |