fixing up py files for run

pom/chords.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 import torch 
 from PIL import Image
 import numpy as np
-from control_toys.utils import rect_to_square, square_to_rect
+from . import rect_to_square, square_to_rect
 
 CHORD_BORDER = 8   # chord border size in pixels
 

pom/utils.py

@@ -0,0 +1,46 @@
+#!/usr/bin/env python3
+
+"""
+simple routines used throughout the project, placed here to avoid circular imports
+"""
+
+from PIL import Image
+import torch 
+
+def flip_bottom_half_and_attach(sub_img):
+    "takes one 256x256 and returns on 512x128 image with the bottom half reversed and attached on the right"
+    h, w = sub_img.size
+    new_img = Image.new(sub_img.mode, (w*2, h//2))
+    new_img.paste(sub_img.crop((0, 0, w, h//2)), (0, 0))
+    new_img.paste(sub_img.crop((0, h//2, w, h)).transpose(Image.FLIP_LEFT_RIGHT), (w, 0))
+    return new_img 
+
+def square_to_rect_tensor(img_tensor:torch.Tensor):
+    if len(img_tensor.shape) <= 4:
+        img_tensor = img_tensor.unsqueeze(0)
+    channels_last = img_tensor.shape[-1] == 3
+    if channels_last:
+        img_tensor = img_tensor.permute(0, 3, 1, 2)
+    channels_first = img_tensor.shape[1] == 3
+    b,c,w,h = img_tensor.shape
+    new_img = torch.zeros((b,c,w*2,h//2), dtype=img_tensor.dtype).to(img_tensor.device)
+    new_img[:,:,:w, :] = img_tensor[:,:,:w,:]
+    new_img[:,:,w:, :] = torch.flip(img_tensor[:,:,w:, :], dims=[-2])
+    return new_img
+
+
+def square_to_rect(img):
+    #"""just an alias for flip_bottom_half_and_attach"""
+    if isinstance(img, torch.Tensor):
+        return square_to_rect_tensor(img)
+    return flip_bottom_half_and_attach(img)
+
+def rect_to_square(img):
+    "takes a 512x128 image and returns a 256x256 image with the bottom half reversed"
+    w, h = img.size
+    new_img = Image.new(img.mode, (w//2, h*2))
+    new_img.paste(img.crop((0, 0, w//2, h)), (0, 0))
+    new_img.paste(img.crop((w//2, 0, w, h)).transpose(Image.FLIP_LEFT_RIGHT), (0, h))
+    return new_img
+
+