import vtk

def addLines(renderer):
    # Create three points. Join (Origin and P0) with a red line and
    # (Origin and P1) with a green line
    origin = [0.0, 77.0, 0.0]
    p0 = [200.0, 77.0, 0.0]
    p1 = [0.0, 0.0, 0.0]

    # Create a vtkPoints object and store the points in it
    pts = vtk.vtkPoints()
    pts.InsertNextPoint(origin)
    pts.InsertNextPoint(p0)
    pts.InsertNextPoint(p1)

    # Setup two colors - one for each line
    red = [255, 0, 0]
    green = [0, 255, 0]

    # Setup the colors array
    colors = vtk.vtkUnsignedCharArray()
    colors.SetNumberOfComponents(3)
    colors.SetName("Colors")

    # Add the colors we created to the colors array
    colors.InsertNextTuple(red)
    colors.InsertNextTuple(green)

    # Create the first line (between Origin and P0)
    line0 = vtk.vtkLine()
    line0.GetPointIds().SetId(0,0) # the second 0 is the index of the Origin in the vtkPoints
    line0.GetPointIds().SetId(1,1) # the second 1 is the index of P0 in the vtkPoints

    # Create the second line (between Origin and P1)
    line1 = vtk.vtkLine()
    line1.GetPointIds().SetId(0,0) # the second 0 is the index of the Origin in the vtkPoints
    line1.GetPointIds().SetId(1,2) # 2 is the index of P1 in the vtkPoints

    # Create a cell array to store the lines in and add the lines to it
    lines = vtk.vtkCellArray()
    lines.InsertNextCell(line0)
    lines.InsertNextCell(line1)

    # Create a polydata to store everything in
    linesPolyData = vtk.vtkPolyData()

    # Add the points to the dataset
    linesPolyData.SetPoints(pts)

    # Add the lines to the dataset
    linesPolyData.SetLines(lines)

    # Color the lines - associate the first component (red) of the
    # colors array with the first component of the cell array (line 0)
    # and the second component (green) of the colors array with the
    # second component of the cell array (line 1)
    linesPolyData.GetCellData().SetScalars(colors)

    # Visualize
    mapper = vtk.vtkPolyDataMapper()
    if vtk.VTK_MAJOR_VERSION <= 5:
        mapper.SetInput(linesPolyData)
    else:
        mapper.SetInputData(linesPolyData)

    actor = vtk.vtkActor()
    actor.SetMapper(mapper)
    renderer.AddActor(actor)

for offscreen in [1, 0]:
    reader = vtk.vtkPNGReader()
    reader.SetFileName('vtk_logo-main1.png')
    reader.Update()

    image = reader.GetOutput()

    dims = image.GetDimensions()
    spacing = image.GetSpacing()

    resliceMapper = vtk.vtkImageResliceMapper()
    resliceMapper.SetInputData(image)

    imageSlice = vtk.vtkImageSlice()
    imageSlice.SetMapper(resliceMapper)

    imageStack = vtk.vtkImageStack()
    imageStack.AddImage(imageSlice)

    renderer = vtk.vtkRenderer()
    renderer.AddActor(imageStack)

    renderer.GetActiveCamera().SetPosition(
        (dims[0]-1)*spacing[0]*0.5,
        (dims[1]-1)*spacing[1]*0.5,
        2000)

    renderer.GetActiveCamera().SetFocalPoint(
        (dims[0]-1)*spacing[0]*0.5,
        (dims[1]-1)*spacing[1]*0.5,
        0)

    renderer.GetActiveCamera().SetParallelProjection(1)
    renderer.GetActiveCamera().SetParallelScale(0.5*dims[1]*spacing[1])
    renderer.ResetCameraClippingRange()

    addLines(renderer)

    renderWindow = vtk.vtkRenderWindow()
    renderWindow.SetOffScreenRendering(offscreen)
    renderWindow.SetSize(200, 78)
    renderWindow.AddRenderer(renderer)

    renderWindow.Render()

    window2image = vtk.vtkWindowToImageFilter()
    window2image.SetInput(renderWindow)

    writer = vtk.vtkPNGWriter()
    writer.SetInputConnection(window2image.GetOutputPort())
    writer.SetFileName(f"screenshot_{offscreen}.png")
    writer.Write()