Update app.py

app.py

@@ -4,6 +4,151 @@ from fastai.vision.all import *
 import PIL
 import torchvision.transforms as transforms
 
+##Extras por si pudiera reconstruir la imagen en HF también
+import os
+import re
+
+def subimages_from_directory(directorio):
+  # Define el directorio a recorrer
+  directorio = directorio
+
+  # Define la expresión regular para buscar los números X e Y en el nombre de archivo
+  patron = re.compile(r"(.*)_(\d+)_(\d+)\.(png|jpg|tif)")
+
+  windowlist = []
+  coords = []
+
+  # Recorre el directorio en busca de imágenes
+  for filename in os.listdir(directorio):
+    match = patron.search(filename)
+    if match:
+        origname = match.group(1)
+        x = int(match.group(2))
+        y = int(match.group(3))
+        #print(f"El archivo {filename} tiene los números X={x} e Y={y}")
+        img = cv2.imread(os.path.join(directorio, filename))
+        windowlist.append(img)
+        coords.append((x, y))
+
+  # Ordena las listas por coordenadas X e Y
+  windowlist, coords = zip(*sorted(zip(windowlist, coords), key=lambda pair: (pair[1][0], pair[1][1])))
+  wh, ww, chan = windowlist[0].shape
+  origsize = tuple(elem1 + elem2 for elem1, elem2 in zip(coords[-1], (wh,ww)))
+
+  return windowlist, coords, wh, ww, chan, origsize
+
+def subimages_onlypath(directorio):
+  # Define el directorio a recorrer
+  directorio = directorio
+  pathlist = []
+
+  patron = re.compile(r"(.*)_(\d+)_(\d+)\.(png|jpg|tif)")
+
+  for filename in os.listdir(directorio):
+    match = patron.search(filename)
+    if match:
+        pathlist.append(os.path.join(directorio, filename))
+
+  return pathlist
+
+def ReconstructFromMW(windowlist, coords, wh, ww, chan, origsize):
+  canvas = np.zeros((origsize[1], origsize[0], chan), dtype=np.uint8)
+  for idx, window in enumerate(windowlist):
+    canvas[coords[idx][1]:coords[idx][1]+wh, coords[idx][0]:coords[idx][0]+ww, :] = window
+  return canvas
+
+def get_list_tp(path):
+  list_to_process = []  # Inicializar la lista que contendrá los nombres de los subdirectorios
+  list_names = []
+  # Recorrer los elementos del directorio
+  for element in os.scandir(path):
+      # Verificar si el elemento es un directorio
+      if element.is_dir():
+          # Agregar el nombre del subdirectorio a la lista
+          windowlist, coords, wh, ww, chan, origsize = subimages_from_directory(element)
+          list_to_process.append(ReconstructFromMW(windowlist, coords, wh, ww, chan, origsize))
+          list_names.append(element.name)
+  return list_to_process, list_names
+
+def get_paths_tp(path):
+  list_to_process = []  # Inicializar la lista que contendrá los nombres de los subdirectorios
+  # Recorrer los elementos del directorio
+  for element in os.scandir(path):
+      # Verificar si el elemento es un directorio
+      if element.is_dir():
+          # Agregar el nombre del subdirectorio a la lista
+          list_to_process.append(subimages_onlypath(element))
+  return list_to_process
+
+def process_multifolder(process_folders, result_folder):
+  for folder in process_folders:
+    folname = os.path.basename(os.path.dirname(folder[0]))
+    destname = Path(result_folder)/folname
+    os.makedirs(destname, exist_ok=True)
+    for subimagepath in folder:
+      img = PIL.Image.open(subimagepath)
+      image = transforms.Resize((400,400))(img)
+      tensor = transform_image(image=image)
+      with torch.no_grad():
+          outputs = model(tensor)
+      outputs = torch.argmax(outputs,1)
+      mask = np.array(outputs.cpu())
+      mask[mask==1]=255
+      mask=np.reshape(mask,(400,400))
+      mask_img = Image.fromarray(mask.astype('uint8'))
+
+      filename = os.path.basename(subimagepath)
+      new_image_path = os.path.join(result_folder, folname, filename)
+      mask_img.save(new_image_path)
+
+def recombine_windows(results_folder_w, result_f_rec):
+  imgs, nombres = get_list_tp(results_folder_w)
+  os.makedirs(result_f_rec, exist_ok=True)
+
+  for idx, image in enumerate(imgs):
+    img = Image.fromarray(image)
+    new_image_path = os.path.join(result_f_rec, nombres[idx] + '.tif')
+    img.save(new_image_path, compression='tiff_lzw')
+  return new_image_path
+
+def process_single_image(single_image_path, base_f, pro_f, rsw_f, rsd_f):
+  gss_single(single_image_path, pro_f, 0, "tif", True)
+  process_multifolder(get_paths_tp(pro_f),rsw_f)
+  pt = recombine_windows(rsw_f,rsd_f)
+  shutil.rmtree(pro_f)
+  shutil.rmtree(rsw_f)
+  copiar_info_georref(single_image_path, pt)
+  return pt
+
+from osgeo import gdal, osr
+
+def copiar_info_georref(entrada, salida):
+    try:
+        # Abrir el archivo GeoTIFF original
+        original_dataset = gdal.Open(entrada)
+
+        # Obtener la información de georreferenciación del archivo original
+        original_projection = original_dataset.GetProjection()
+        original_geotransform = original_dataset.GetGeoTransform()
+
+        # Abrir la imagen resultado
+        result_dataset = gdal.Open(salida, gdal.GA_Update)
+
+        # Copiar la información de georreferenciación del archivo original a la imagen resultado
+        result_dataset.SetProjection(original_projection)
+        result_dataset.SetGeoTransform(original_geotransform)
+
+        # Cerrar los archivos
+        original_dataset = None
+        result_dataset = None
+
+    except Exception as e:
+        print("Error: ", e)
+
+###FIN de extras
+
+
+
 #repo_id = "Ignaciobfp/segmentacion-dron-marras"
 #learner = from_pretrained_fastai(repo_id)