PERFEKTROTATOR for all your rotation needs

rotator.py

import sys
import os
import os.path
import math
from collections import namedtuple
# screw this
os.environ["OPENCV_IO_MAX_IMAGE_PIXELS"] = str(pow(2,40))
import cv2 as cv

PointBase = namedtuple("Point", ("x", "y", "z"), defaults=(0, 0, 0))
class Point(PointBase):
  def __mul__(self, scalar):
    return Point(self.x * scalar, self.y * scalar, self.z * scalar)
  def __truediv__(self, scalar):
    return Point(self.x / scalar, self.y / scalar, self.z / scalar)
  def __add__(self, other):
    if isinstance(other, Vector):
      return Point(self.x + other.x, self.y + other.y, self.z + other.z)
    return NotImplemented
  def __sub__(self, other):
    if isinstance(other, Vector):
      return Point(self.x - other.x, self.y - other.y, self.z - other.z)
    if isinstance(other, Point):
      return Vector(self.x - other.x, self.y - other.y, self.z - other.z)
    return NotImplemented
  def __eq__(self, other):
    if isinstance(other, Point):
      return self.x == other.x and self.y == other.y and self.z == other.z
    return NotImplemented
  def __hash__(self):
    return hash((self.x, self.y, self.z))
  @property
  def xy(self):
    return (self.x, self.y)

VectorBase = namedtuple("Vector", ("x", "y", "z"), defaults=(0, 0, 0))
class Vector(VectorBase):
  def __mul__(self, scalar):
    return Vector(self.x * scalar, self.y * scalar, self.z * scalar)
  def __truediv__(self, scalar):
    return Vector(self.x / scalar, self.y / scalar, self.z / scalar)
  @property
  def length(self):
    return pow(self.x*self.x + self.y*self.y + self.z*self.z, 0.5)
  def __neg__(self):
    return Vector(-self.x, -self.y, -self.z)
  def __add__(self, other):
    if isinstance(other, Vector):
      return Vector(self.x + other.x, self.y + other.y, self.z + other.z)
    if isinstance(other, tuple):
      if len(other) == 3:
        return Vector(self.x + other[0], self.y + other[1], self.z + other[2])
      if len(other) == 2:
        return Vector(self.x + other[0], self.y + other[1], self.z)
    return NotImplemented
  def __sub__(self, other):
    if isinstance(other, Vector):
      return Vector(self.x - other.x, self.y - other.y, self.z - other.z)
    if isinstance(other, tuple):
      if len(other) == 3:
        return Vector(self.x - other[0], self.y - other[1], self.z - other[2])
      if len(other) == 2:
        return Vector(self.x - other[0], self.y - other[1], self.z)
    return NotImplemented
  def __eq__(self, other):
    if isinstance(other, Vector):
      return self.x == other.x and self.y == other.y and self.z == other.z
    return NotImplemented
  def __hash__(self):
    return hash((self.x, self.y, self.z))
  @property
  def xy(self):
    return (self.x, self.y)

RectBase = namedtuple("Rect", ("top_left", "bottom_right"))
class Rect(RectBase):
  @staticmethod
  def from_xywh(top_left, extents):
    return Rect(top_left, top_left + extents)
  @property
  def top_right(self):
    return Point(self.bottom_right.x, self.top_left.y)
  @property 
  def bottom_left(self):
    return Point(self.top_left.x, self.bottom_right.y)
  @property
  def width(self):
    return self.bottom_right.x - self.top_left.x
  @property
  def height(self):
    return self.bottom_right.y - self.top_left.y
  @property
  def center(self):
    return self.top_left + Vector(0.5 * self.width, 0.5 * self.height)
  def __add__(self, other):
    if isinstance(other, Vector):
      return Rect(self.top_left + other, self.bottom_right + other)
    return NotImplemented
  def __hash__(self):
    return hash((self.top_left, self.bottom_right))
  def contains(self, point):
    return self.top_left.x <= point.x <= self.bottom_right.x and self.top_left.y <= point.y <= self.bottom_right.y
  def overlaps(self, other):
    if (self.width + other.width) < abs(self.center.x - other.center.x):
      return False
    if (self.height + other.height) < abs(self.center.y - other.center.y):
      return False
    return True


class MainWindow:
    """
    This window provides the overview over the current state of affairs and
    allows editing the project data.
    """

    WINDOW_NAME = "PERFEKTROTATOR"

    def __init__(self, src_image_path, dst_image_path=None):
        print("loading image ... will take a while ...")
        self.image = cv.imread(src_image_path)
        self.dst_image_path = dst_image_path
        print("pick two points on a horizontal line and press enter to rotate and save.")
        print("f1/f2 zooms, 1/2 chooses the point (starts with 1st), mouse lmb picks.")
        # this is the viewed space
        image_height, image_width, _ = self.image.shape
        self._view_center = Vector(image_width/2, image_height/2)
        self._view_scale = 1.0 # 2 = two screen pixels per image pixel
        # this is the actual render window size (width, height), in screen coords
        self._viewport_size = None
        self._redraw_needed = False
        # points 1 and 2
        self.points = [None, None]
        self._choose_point(0)
        # do the opencv horrors
        cv.namedWindow(self.WINDOW_NAME, cv.WINDOW_GUI_EXPANDED)
        cv.setMouseCallback(self.WINDOW_NAME, self._mouse_callback)
        # ready~
        #self._adjust_view_center()

    def translate_view(self, offset):
        self._view_center += offset
        # make sure the view doesn't fall outside the image
        self._adjust_view_center()
        # flag for redraw automatically
        self._redraw_needed = True

    def scale_view(self, scale):
        self._view_scale = self._view_scale * scale
        if self._view_scale > 3.0:
            self._view_scale = 3.0
        # force width/height to max image size by adjusting the scale
        image_height, image_width, _ = self.image.shape
        view_area = self._viewport_size / self._view_scale
        if view_area.x > image_width:
            self._view_scale = self._viewport_size.x / image_width
        if view_area.y > image_height:
            self._view_scale = max(self._view_scale, self._viewport_size.y / image_height)
        # make sure the view doesn't fall outside the image
        self._adjust_view_center()
        # flag for redraw automatically
        self._redraw_needed = True

    def _adjust_view_center(self):
        # note: scale is already constrained
        image_height, image_width, _ = self.image.shape
        view_area = self._viewport_size / self._view_scale
        top_left = self._view_center - view_area * 0.5
        self._view_center += -Vector(min(top_left.x, 0), min(top_left.y, 0))
        btm_right = self._view_center + view_area * 0.5
        self._view_center += Vector(min(image_width - btm_right.x, 0), min(image_height - btm_right.y, 0))

    def _screen_to_image(self, point):
        # screen coords are relative to top left
        return (point - self._viewport_size * 0.5) / self._view_scale + self._view_center
    def _image_to_screen(self, point):
        point = (point - self._view_center) * self._view_scale + self._viewport_size * 0.5
        return Vector(int(point.x), int(point.y))

    def _redraw(self):
        # this occurs when a window is minimized; don't draw anything
        if self._viewport_size.x == 0 or self._viewport_size.y == 0:
            return
        # compute the source area
        view_area = self._viewport_size / self._view_scale
        view_rect = Rect(self._view_center - view_area * 0.5, self._view_center + view_area * 0.5)
        # crop out the viewed area; 
        image = self.image[int(view_rect.top_left.y):int(view_rect.bottom_right.y), int(view_rect.top_left.x):int(view_rect.bottom_right.x)]
        # rescale the image; this is now our background
        background = cv.resize(image, (self._viewport_size.x, self._viewport_size.y), interpolation=cv.INTER_LINEAR)
        # submit to opencv for display
        cv.imshow(self.WINDOW_NAME, background)

    def _mouse_callback(self, event, x, y, flags, _):
        # flag for redraw automatically
        self._redraw_needed = True
        # handle scaling
        amount = flags / 131072
        if event == cv.EVENT_MOUSEWHEEL:
            self.translate_view(Vector(0, -amount) / self._view_scale)
        elif event == cv.EVENT_MOUSEHWHEEL:
            self.translate_view(Vector(amount, 0) / self._view_scale)
        elif event == cv.EVENT_LBUTTONDOWN:
            point = self._screen_to_image(Vector(x, y))
            print("Point %d at (%.1f, %.1f)" % (self.point_index+1, point.x, point.y))
            self.points[self.point_index] = point
            if self.point_index == 0:
                self._choose_point(1)

    def _choose_point(self, index):
        self.point_index = index
        print("choosing point %d" % (index + 1))

    def _handle_key(self, keycode):
        # flag for redraw automatically
        self._redraw_needed = True

        if keycode == 0x700000: # f1
            self._view_scale *= 2
            self._adjust_view_center()
        elif keycode == 0x710000: # f2
            self._view_scale /= 2
            self._adjust_view_center()
        elif keycode == ord('1'):
            self._choose_point(0)
        elif keycode == ord('2'):
            self._choose_point(1)
        elif keycode == 0x0d: # enter
            if self.points[0] is not None and self.points[1] is not None:
                angle = math.atan2(self.points[1].y - self.points[0].y, self.points[1].x - self.points[0].x) * 180 / math.pi
                print("estimated rotation angle: %.2f deg" % (angle,))
                if self.dst_image_path is not None:
                    print("rotating ...")
                    height, width = self.image.shape[:2]
                    rotate_matrix = cv.getRotationMatrix2D(center=(width/2, height/2), angle=angle, scale=1)
                    image = cv.warpAffine(src=self.image, M=rotate_matrix, dsize=(width, height))
                    print("saving ...")
                    cv.imwrite(self.dst_image_path, image)
                    print("done.")
                    sys.exit()

    def main_loop(self):
        while True:
            # check if still running
            window_is_open = cv.getWindowProperty(self.WINDOW_NAME, cv.WND_PROP_VISIBLE)
            if not window_is_open:
                sys.exit()

            # handle window resizing
            viewport = cv.getWindowImageRect(self.WINDOW_NAME)[2:]
            viewport = Vector(viewport[0], viewport[1])
            if viewport != self._viewport_size:
                #print("resized; new size: %dx%d" % (viewport.x, viewport.y))
                self._viewport_size = viewport
                self._adjust_view_center()
                self._redraw_needed = True

            # handle inputs
            keycode = cv.waitKeyEx(1)
            if keycode != -1:
                self._handle_key(keycode)

            # handle redraws
            if self._redraw_needed:
                self._redraw_needed = False
                self._redraw()

if __name__ == "__main__":
    print("PERFEKTROTATOR by @InfoSecDJ, winter 2023. no version number needed.")
    if len(sys.argv) < 2:
        print("use: %s <path to source image> [maybe path to dest image]")
        sys.exit(1)
    try:
        dst_image_path = sys.argv[2]
    except:
        dst_image_path = None
    win = MainWindow(sys.argv[1], dst_image_path)
    win.main_loop()